import matplotlib.pyplot as plt
import numpy as np

plt.rcParams["font.family"] = ["SimHei"]
import pandas as pd
from scipy.ndimage import gaussian_filter

# 导入自定义模块
from initialization import initialization
from pso import pso
from aimfcn_1 import aimFcn_1
from draw_pc import draw_pc


def main():
    result_dict = {}  # 初始化result_dict以避免UnboundLocalError
    """3D路径规划主程序"""
    plt.close('all')  # 关闭所有现有图形
    np.random.seed(42)  # 设置随机种子，保证结果可复现
    
    # -------------------------- 数据加载与环境配置 ------------------------- #
    # 加载地图数据
    data = {}
    data['map'] = pd.read_excel('map1.xlsx', header=None).values  # 读取Excel地图数据
    data['minF'] = np.array([10, 0, 0])
    data['maxF'] = np.array([100, 1, 100])
    data['mu_max'] = np.pi / 3  # 最大俯仰角
    data['beta_max'] = np.pi / 3  # 最大转弯角
    
    # 构建三维地图
    data['mapsize'] = [100, 100]
    data['mapSize0'] = data['mapsize']  # 添加mapSize0键
    # 从Excel加载地图数据
    df = pd.read_excel('map1.xlsx', header=None)
    data['map_z'] = df.values  # 地形高度数据
    t = np.arange(1, 101)
    x_grid, y_grid = np.meshgrid(t, t)  # 创建网格
    
    # 基准地形建模
    h1 = (np.sin(y_grid + 10) + 0.2 * np.sin(x_grid) + 0.1 * np.cos(0.6 * np.sqrt(x_grid**2 + y_grid**2)) +
          np.cos(y_grid) + np.sin(0.1 * np.sqrt(x_grid**2 + y_grid**2)) + 0.1 * np.cos(y_grid))
    h1 = (h1 + 2 * np.random.randn(*h1.shape)) * 10
    h1 = gaussian_filter(h1, sigma=3)  # 高斯滤波平滑地形
    
    # 山峰建模、威胁区域建模
    data['Ob'] = np.array([
        [40, 20, 5, 15, 10], [20, 40, 5, 6, 5], [60, 80, 5, 10, 10],
        [80, 60, 5, 6, 5], [20, 50, 5, 10, 10], [70, 20, 5, 6, 5], [80, 20, 5, 6, 5]
    ])
    data['numOb'] = data['Ob'].shape[0]
    h_peak = np.array([28, 21, 26, 30, 30, 19, 20])  # 威胁区域高度
    
    # 生成山峰地形
    h2 = np.zeros_like(h1)
    for i in range(data['numOb']):
        y0, x0 = data['Ob'][i, :2]
        h = h_peak[i]
        xs, ys = data['Ob'][i, 3:5]
        h2 += h * np.exp(-((x_grid - x0)/xs)**2 - ((y_grid - y0)/ys)**2)
    
    # 合并地形
    data['map_z'] = np.maximum(h1, h2)  # 最终地形高度
    data['map_x'], data['map_y'] = x_grid, y_grid
    
    # 路径规划参数设置
    data['minH'] = 10  # 最低飞行高度
    data['maxH'] = 20  # 最高飞行高度
    data['S'] = np.array([0, 0, 10])  # 设置起点坐标
    data['noS'] = 0  # 起点索引
    data['E0'] = np.array([[100, 100, 15]])  # 终点坐标
    data['noE0'] = [9999]  # 恢复为列表格式以控制路径数量
    data['noP'] = 0  # 路径索引参数
    
    # 节点网络构建
    data['node'] = np.column_stack((
        np.repeat(np.arange(1, 101), 100),
        np.tile(np.arange(1, 101), 100),
        data['map_z'].flatten()
    ))
    data['numNode'] = data['node'].shape[0]

    # 构建节点邻接网络 (4邻域连接)
    net = []
    mapsize = data['mapsize']
    for i in range(mapsize[0]):
        for j in range(mapsize[1]):
            current_idx = i * mapsize[1] + j
            # 四向邻接连接
            # 8邻域连接（4方向+4对角线）
            # 右邻接
            if j < mapsize[1] - 1:
                right_idx = i * mapsize[1] + (j + 1)
                net.append([current_idx, right_idx])
            # 左邻接
            if j > 0:
                left_idx = i * mapsize[1] + (j - 1)
                net.append([current_idx, left_idx])
            # 下邻接
            if i < mapsize[0] - 1:
                down_idx = (i + 1) * mapsize[1] + j
                net.append([current_idx, down_idx])
            # 上邻接
            if i > 0:
                up_idx = (i - 1) * mapsize[1] + j
                net.append([current_idx, up_idx])
            # 右下邻接
            if i < mapsize[0] - 1 and j < mapsize[1] - 1:
                down_right_idx = (i + 1) * mapsize[1] + (j + 1)
                net.append([current_idx, down_right_idx])
            # 左下邻接
            if i < mapsize[0] - 1 and j > 0:
                down_left_idx = (i + 1) * mapsize[1] + (j - 1)
                net.append([current_idx, down_left_idx])
            # 右上邻接
            if i > 0 and j < mapsize[1] - 1:
                up_right_idx = (i - 1) * mapsize[1] + (j + 1)
                net.append([current_idx, up_right_idx])
            # 左上邻接
            if i > 0 and j > 0:
                up_left_idx = (i - 1) * mapsize[1] + (j - 1)
                net.append([current_idx, up_left_idx])
    data['net'] = np.array(net, dtype=int)

    # -------------------------- 算法参数配置 ------------------------- #
    option = {}
    option['dim'] = len(data['E0'])  # 大幅降低维度，仅保留终点排序参数
    option['lb'] = 0.0  # 进一步扩大搜索空间
    option['ub'] = 3.0  # 进一步扩大搜索空间
    option['fobj'] = aimFcn_1  # 目标函数
    
    # 统一边界长度
    if np.isscalar(option['lb']):
        option['lb'] = np.ones(option['dim']) * option['lb']
        option['ub'] = np.ones(option['dim']) * option['ub']
    
    # 算法参数
    option['numAgent'] = 25  # 优化种群规模平衡多样性与计算效率
    option['maxIteration'] = 45   # 增加迭代次数提供充足搜索时间
    option['w_pso'] = 0.9  # 初始惯性权重，将在PSO中自适应调整
    # 移除固定加速系数，使用PSO内部的时变加速策略
    option['stagnation_threshold'] = 12  # 优化停滞阈值平衡探索与收敛
    option['early_stop_threshold'] = 15  # 降低早停阈值减少无效迭代
    option['epsilon'] = 1e-5  # 放宽改进阈值，避免过早停止
    
    # -------------------------- 初始化与求解 ------------------------- #
    # 初始化种群
    x = initialization(option['numAgent'], option['dim'], option['ub'], option['lb'])
    y = np.array([aimFcn_1(xi, option, data)['fit'] for xi in x])
    
    # 使用PSO算法求解
    print("开始PSO算法优化...")
    bestY, bestX, recording = pso(x, y, option, data)
    print(f"优化完成！最优适应度值: {bestY:.4f}")
    
    # -------------------------- 结果可视化 ------------------------- #

    # 绘制适应度曲线
    plt.figure(figsize=(10, 4))
    plt.subplot(121)
    plt.plot(recording['bestFit'], 'b-', linewidth=2)
    plt.xlabel('迭代次数')
    plt.ylabel('适应度值')
    plt.title('适应度函数曲线')
    plt.grid(True)
    
    plt.subplot(122)
    plt.plot(np.log(recording['bestFit']), 'r-', linewidth=2)
    plt.xlabel('迭代次数')
    plt.ylabel('log(适应度值)')
    plt.title('适应度函数曲线(log)')
    plt.grid(True)
    plt.tight_layout()
    
    # 绘制优化路径
    # 调用目标函数计算最优路径并获取结果
    result = aimFcn_1(bestX, option, data)
    # 从结果中提取优化路径
    pathR = result['path']
    # 将路径转换为numpy数组格式
    path_coords = np.concatenate(pathR)[:, 1:4]  # 提取x,y,z坐标
    # 验证路径数据
    print('优化路径点数量:', path_coords.shape[0])
    print('路径起点:', path_coords[0])
    print('路径终点:', path_coords[-1])
    # 保留numpy数组格式以支持draw_pc中的数值操作
    # 最终确认路径数据格式
    # 将路径数组包装为列表以支持路径索引
    # 将路径数组转换为纯Python列表格式以支持正确索引
    path_list = path_coords.tolist()
    # 确保路径为numpy数组格式以支持shape属性
    # 修正路径数据结构为直接numpy数组而非列表嵌套
    # 恢复路径数据结构为列表嵌套以支持多路径索引
    # 验证路径数据并添加调试输出
    print('路径数据形状:', np.array(path_list).shape)
    print('路径数据前5个点:', np.array(path_list)[:])
    result_dict = {'path': [np.array(path_list)], 'fit': bestY}  # 路径列表包含单个numpy数组

    data['noP'] = 0  # 路径索引参数
    newPath = draw_pc(result_dict, option, data, 'PSO')
    
    plt.show()
    return newPath


if __name__ == '__main__':
    # 执行主程序
    newPath = main()
    print("程序执行完成！")